"""
    1、使用gdal_array将图片作为一个数组载入
    2、使用PyShp库创建一个shapefile 文件读取器
    3、栅格化shapefile文件使之成为一张包含地理参照信息的图片（将矢量转换为栅格）
    4、将Shapefile文件影像作为过滤器只获取Shapefile文件边框的影像像素值
    5、使用上述遮罩图片对卫星影像过滤
    6、删除不在边框范围内的卫星影像数据
    7、将裁剪过的卫星影像另存为 clip.tif
"""
import operator
from osgeo import gdal, gdal_array, osr
import shapefile
from PIL import Image, ImageDraw


# 用于剪裁的栅格影像
tifPath = "TIF/FalseColor/stretched.tif"
# 用于裁剪栅格文件的过变形shp文件
shpPath = "TIF/FalseColor/hancock/hancock"
# 裁剪后的栅格文件名
output = "TIF/FalseColor/hancock/clip"


# 将一个Python影像库的数组转换为一个gdal_array图片
def imageToArray(i):
    a = gdal_array.numpy.frombuffer(i.tobytes(), "b")
    a.shape = i.im.size[1], i.im.size[0]
    return a


# 使用GDAL库的geomatrix对象计算地理坐标的像素位置
def world2Pixel(geoMatrix, x, y):
    ulX = geoMatrix[0]
    ulY = geoMatrix[3]
    xDist = geoMatrix[1]
    yDist = geoMatrix[5]
    rtnX = geoMatrix[2]
    rtnY = geoMatrix[4]
    pixel = int((x - float(ulX)) / float(xDist))
    line = int((float(ulY) - y) / abs(float(yDist)))
    return pixel, line


# 将源数据作为gdal_array载入
srcArray = gdal_array.LoadFile(tifPath)

# 同时载入gdal库的图片从而获取 geotransform（世界文件）
srcImage = gdal.Open(tifPath)
geoTrans = srcImage.GetGeoTransform()

# 使用PyShp库打开shp文件
r = shapefile.Reader(shpPath)

# 将图层扩展转换为图片像素坐标
minX, minY, maxX, maxY = r.bbox
ulX, ulY = world2Pixel(geoTrans, minX, maxY)
lrX, lrY = world2Pixel(geoTrans, maxX, minY)

# 就是那图片的像素尺寸
pxWidth = int(lrX - ulX)
pxHeight = int(lrY - ulY)
clip = srcArray[:, ulY:lrY, ulX:lrX]

# 为图片创建一个新的 geomatrix 对象一边附加地理参照数据
geoTrans = list(geoTrans)
geoTrans[0] = minX
geoTrans[3] = maxY

# 边界线
pixels = []
for p in r.shape(0).points:
    pixels.append(world2Pixel(geoTrans, p[0], p[1]))
rasterPoly = Image.new("L", (pxWidth, pxHeight), 1)

# 使用PLI创建一个空白图片用于绘制多边形
rasterize = ImageDraw.Draw(rasterPoly)
rasterize.polygon(pixels, 0)

# 将PIL图片转换为NumPy遮罩数组
mask = imageToArray(rasterPoly)

# 根据遮罩图层对图片进行裁剪
clip = gdal_array.numpy.choose(mask, (clip, 0)).astype(gdal_array.numpy.uint8)
print(clip)
# 将 ndvi 保存为tiff文件
gdal_array.SaveArray(clip, "{}.tif".format(output),
                     format="GTiff", prototype=tifPath)
